import os
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional
import logging

logger = logging.getLogger(__name__)

def analyze_reboot_times_and_assign_stages(logs: List[Dict]) -> Dict[str, Any]:
    """
    分析日志中的重启时间，并为每个日志条目分配运行阶段ID。
    该函数从给定的日志列表中检测重启事件，并为每个日志条目分配一个
    唯一的运行阶段ID。当检测到重启时，run_stage_id 会递增。

    Args:
        logs: 原始日志条目列表，每个条目是一个字典，至少包含 'datetime' 和 'message'。

    Returns:
        Dict[str, Any]: 包含 'reboot_events' 列表和 'fixed_logs' 列表的字典。
                        'fixed_logs' 中的每个日志条目都将包含 'run_stage_id'。
    """
    reboot_events = []
    processed_logs_with_stages = [] # 存储所有带有阶段ID的日志
    
    current_run_stage_id = 0
    reboot_start_log_dt: Optional[datetime] = None # 记录当前重启阶段开始的时间（在日志中观察到的时间）
    reboot_event_index: Optional[int] = None # 记录reboot_events中当前重启事件的索引
    
    last_processed_log_dt: Optional[datetime] = None # 记录上一个已处理日志的datetime

    for i, log in enumerate(logs):
        log_copy = log.copy() # 处理日志副本

        try:
            log_dt = datetime.strptime(log_copy['datetime'], '%Y-%m-%d %H:%M:%S.%f')
        except ValueError:
            try:
                log_dt = datetime.strptime(log_copy['datetime'], '%Y-%m-%d %H:%M:%S')
            except ValueError:
                log_copy['run_stage_id'] = current_run_stage_id
                processed_logs_with_stages.append(log_copy)
                continue

        is_1970 = log_dt.year == 1970
        
        is_reboot_triggered = False

        # 1. 检测1970年时间跳变或特定日志触发重启
        if (is_1970 and (last_processed_log_dt is None or last_processed_log_dt.year >= 2000)) or \
           ("[:00000002] framework: LAUNCH snlua bootstrap" in log_copy['message']):
            
            # 只有当不在重启期间时才视为新的重启开始
            if reboot_start_log_dt is None:
                current_run_stage_id += 1
                reboot_start_log_dt = log_dt
                reboot_events.append({
                    'reboot_start_log_time': reboot_start_log_dt.isoformat(timespec='microseconds'),
                    'reboot_end_log_time': None, # 暂时未知
                    'estimated_real_reboot_time': None, # 暂时未知
                    'run_stage_id': current_run_stage_id # 记录重启事件的阶段ID
                })
                reboot_event_index = len(reboot_events) - 1
                logger.info(f"检测到重启开始，阶段ID: {current_run_stage_id}: {log_dt}")
            is_reboot_triggered = True

        # 2. 检测1970年内部时间回跳
        if is_1970 and last_processed_log_dt and last_processed_log_dt.year == 1970 and log_dt < last_processed_log_dt:
            # 即使已经在1970期间，内部回跳也视为新的重启
            current_run_stage_id += 1
            reboot_start_log_dt = log_dt # 更新重启开始时间
            reboot_events.append({
                'reboot_start_log_time': reboot_start_log_dt.isoformat(timespec='microseconds'),
                'reboot_end_log_time': None,
                'estimated_real_reboot_time': None,
                'run_stage_id': current_run_stage_id
            })
            reboot_event_index = len(reboot_events) - 1
            logger.info(f"检测到1970年内部时间回跳，标记为新的重启开始，阶段ID: {current_run_stage_id}: {log_dt}")
            is_reboot_triggered = True


        # 3. 处理从1970年到正常年份的跳变 (重启结束)
        if reboot_start_log_dt is not None and not is_1970 and log_dt.year >= 2000:
            reboot_end_log_dt = log_dt
            
            # 在当前重启阶段中找到最后一个1970年日志
            last_1970_log_dt_in_period: Optional[datetime] = None
            for prev_log in reversed(processed_logs_with_stages):
                if prev_log.get('run_stage_id') == current_run_stage_id and prev_log['datetime'].startswith('1970'):
                    try:
                        last_1970_log_dt_in_period = datetime.strptime(prev_log['datetime'], '%Y-%m-%d %H:%M:%S.%f')
                    except ValueError:
                        try:
                            last_1970_log_dt_in_period = datetime.strptime(prev_log['datetime'], '%Y-%m-%d %H:%M:%S')
                        except ValueError:
                            pass
                    if last_1970_log_dt_in_period and last_1970_log_dt_in_period.year == 1970:
                        break
                elif prev_log.get('run_stage_id') != current_run_stage_id:
                    # 如果追溯到上一个阶段，则停止
                    break

            if last_1970_log_dt_in_period and reboot_end_log_dt:
                time_in_1970_era = (last_1970_log_dt_in_period - reboot_start_log_dt).total_seconds()
                # 调整推断逻辑，确保它与原始的LogProcessor逻辑一致
                inferred_real_start_time = reboot_end_log_dt - timedelta(seconds=time_in_1970_era + 1)
                time_offset = inferred_real_start_time - reboot_start_log_dt

                # 更新reboot_events
                if reboot_event_index is not None:
                    reboot_events[reboot_event_index].update({
                        'reboot_end_log_time': reboot_end_log_dt.isoformat(timespec='microseconds'),
                        'estimated_real_reboot_time': (reboot_end_log_dt - (last_1970_log_dt_in_period - reboot_start_log_dt)).isoformat(timespec='microseconds')
                    })
                
                # 应用时间修复
                for p_log in processed_logs_with_stages:
                    if p_log.get('run_stage_id') == current_run_stage_id:
                        p_log_dt_str = p_log['datetime']
                        try:
                            p_log_dt = datetime.strptime(p_log_dt_str, '%Y-%m-%d %H:%M:%S.%f')
                        except ValueError:
                            try:
                                p_log_dt = datetime.strptime(p_log_dt_str, '%Y-%m-%d %H:%M:%S')
                            except ValueError:
                                continue # 跳过无法解析时间的日志

                        if p_log_dt.year == 1970:
                            fixed_dt = p_log_dt + time_offset
                            p_log['datetime'] = fixed_dt.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
                            p_log['is_fixed_time'] = True
                logger.info(f"检测到重启结束并修复时间，阶段ID: {current_run_stage_id}: {reboot_end_log_dt}")
            
            # 重置重启状态
            reboot_start_log_dt = None
            reboot_event_index = None
        
        # 分配阶段ID
        log_copy['run_stage_id'] = current_run_stage_id
        processed_logs_with_stages.append(log_copy)
        last_processed_log_dt = log_dt

    return {
        'reboot_events': reboot_events,
        'fixed_logs': processed_logs_with_stages
    } 